from griddata import griddata
import matplotlib.pyplot as plt
import numpy as np

import scikits.delaunay as delaunay
from numpy import ma
def griddata2(x,y,z,xi,yi):
    """
    zi = griddata(x,y,z,xi,yi) fits a surface of the form z = f(x,y)
    to the data in the (usually) nonuniformly spaced vectors (x,y,z).
    griddata interpolates this surface at the points specified by (xi,yi)
    to produce zi. xi and yi must describe a regular grid.
    
    Uses natural neighbor interpolation based on Delaunay triangulation.
    """
    if xi.ndim != yi.ndim:
        raise TypeError("inputs xi and yi must have same number of dimensions (1 or 2)")
    if xi.ndim != 1 and xi.ndim != 2:
        raise TypeError("inputs xi and yi must be 1D or 2D.")
    if xi.ndim == 1:
        xi,yi = np.meshgrid(xi,yi)
    # triangulate data
    tri = delaunay.Triangulation(x,y)
    # interpolate data
    interp = tri.nn_interpolator(z)
    zi = interp(xi,yi)
    # mask points on grid outside convex hull of input data.
    zi = ma.masked_where(np.isnan(zi),zi)
    return zi
#griddata = griddata2

# test case that scikits.delaunay fails on.
data = np.array([[-1, -1], [-1, 0], [-1, 1],
                 [ 0, -1], [ 0, 0], [ 0, 1],
                 [ 1, -1 - np.finfo(np.float_).eps], [ 1, 0], [ 1, 1],
                ])
x = data[:,0]
y = data[:,1]
z = x*np.exp(-x**2-y**2)
# define grid.
xi = np.linspace(-1.1,1.1,100)
yi = np.linspace(-1.1,1.1,100)
# grid the data.
zi = griddata(x,y,z,xi,yi)
# contour the gridded data, plotting dots at the nonuniform data points.
CS = plt.contour(xi,yi,zi,15,linewidths=0.5,colors='k')
CS = plt.contourf(xi,yi,zi,15,cmap=plt.cm.jet)
plt.colorbar() # draw colorbar
# plot data points.
plt.scatter(x,y,marker='o',c='b',s=5)
plt.xlim(-1.1,1.1)    
plt.ylim(-1.1,1.1)     
plt.title('griddata test 2')
plt.show()
